import torch.nn as nn
import torch
from functions import ReverseLayerF


class CNNModel(nn.Module):

    def __init__(self,hparams):
        super(CNNModel, self).__init__()
        self.feature_conv = nn.Sequential()

        self.feature_conv.add_module('f_conv1', nn.Conv2d(1, 16, kernel_size=(1, 7), padding='same'))
        self.feature_conv.add_module('f_bn1', nn.BatchNorm2d(16))
        self.feature_conv.add_module('f_pool1', nn.MaxPool2d(kernel_size=(1, 2)))
        self.feature_conv.add_module('f_relu1', nn.ReLU(True))
        self.feature_conv.add_module('f_conv2',nn.Conv2d(16, 24, kernel_size=(2, 7), padding='same'))
        self.feature_conv.add_module('f_bn2', nn.BatchNorm2d(24))
        self.feature_conv.add_module('f_drop1', nn.Dropout2d())
        self.feature_conv.add_module('f_pool2', nn.MaxPool2d(kernel_size=(2, 2)))
        self.feature_conv.add_module('f_relu2', nn.ReLU(True))
        # self.feature_conv.add_module('f_conv3',nn.Conv2d(32, 64, kernel_size=(1, 7), padding='same'))
        # self.feature_conv.add_module('f_bn3', nn.BatchNorm2d(64))
        # self.feature_conv.add_module('f_drop1', nn.Dropout2d())
        # self.feature_conv.add_module('f_pool3', nn.MaxPool2d(kernel_size=(1, 2)))
        # self.feature_conv.add_module('f_relu3', nn.ReLU(True))

        self.feature_fc = nn.Sequential()
        self.feature_fc.add_module('f_fc1', nn.Linear(int(9600/2), 400)) #600
        # self.feature_fc.add_module('f_fc1', nn.Linear(4768, 400))
        self.feature_fc.add_module('f_fc1', nn.Linear(768, 400))
        ## ----------------标签预测模型部分--------------------
        self.class_classifier = nn.Sequential()
        self.class_classifier.add_module('c_fc1', nn.Linear(400,64))
        self.class_classifier.add_module('c_bn1', nn.BatchNorm1d(64))
        self.class_classifier.add_module('c_relu1', nn.ReLU(True))
        self.class_classifier.add_module('c_drop1', nn.Dropout())
        self.class_classifier.add_module('c_fc2', nn.Linear(64, 1))
        # self.class_classifier.add_module('c_bn2', nn.BatchNorm1d(100))
        # self.class_classifier.add_module('c_relu2', nn.ReLU(True))
        # self.class_classifier.add_module('c_fc3', nn.Linear(100, 1))
       # self.class_classifier.add_module('c_softmax', nn.LogSoftmax(dim=1))
        self.class_classifier.add_module('c_tanh', nn.Tanh())
        # ---------------------------------------------------

        ## ---------------域分类预测模型部分-------------------
        self.domain_classifier = nn.Sequential()
        self.domain_classifier.add_module('d_fc1', nn.Linear(400, 64))
        self.domain_classifier.add_module('d_bn1', nn.BatchNorm1d(64))
        self.domain_classifier.add_module('d_relu1', nn.ReLU(True))
        self.domain_classifier.add_module('d_fc2', nn.Linear(64, 1))
        #self.class_classifier.add_module('d_tanh', nn.Tanh())
        self.domain_classifier.add_module('d_softmax', nn.LogSoftmax(dim=1))
        # ----------------------------------------------------

    def _forward(self,input):
        y=self.feature_conv(input)
        y = y.view(y.shape[0], -1)
        # print(y.shape)
        feature=self.feature_fc(y)
        return feature

    def forward(self, input_x1,input_x2, alpha):
        # input_data = input_data.expand(input_data.data.shape[0], 3, 28, 28)
        feature_x1 = self._forward(input_x1)
        feature_x2 = self._forward(input_x2)
        feature_pair = torch.abs(feature_x2 - feature_x1)
        reverse_feature_pair = ReverseLayerF.apply(feature_pair, alpha)
        class_output = self.class_classifier(feature_pair)
        domain_output = self.domain_classifier(reverse_feature_pair)

        return class_output, domain_output,feature_x1,feature_x2

